import pygrib
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import config


def read_file(filename):
    grbs = pygrib.open(filename)
    grb = grbs.read()[0]
    return grb


def print_all_key(grb):
    keys = grb._all_keys
    for key in keys:
        print(key, grb[key])


def lon_lat_bounds(grb):
    lat, lon = grb.latlons()

    lon_bounds = (lon.min(), lon.max())  # 经度longitude
    lat_bounds = (lat.min(), lat.max())  # 纬度latitude

    # print("经度lon:", lon_bounds, "纬度lat:", lat_bounds)

    return (lon_bounds, lat_bounds)


def show_hot(data):
    # cmap = matplotlib.cm.jet
    # # 设定每个图的colormap和colorbar所表示范围是一样的，即归一化
    # norm = matplotlib.colors.Normalize(vmin=0, vmax=1.2)
    # # 显示图形，此处没有使用contourf #>>>ctf=plt.contourf(grid_x,grid_y,grid_z)
    # gci = plt.imshow(data, cmap=cmap, norm=norm)
    # # gci = plt.imshow(data)
    #
    # # ax = plt.gca()
    # # ax.set_yticks(np.linspace(0, 10, 8))
    # # ax.set_yticklabels(('0.60', '0.65', '0.70', '0.75', '0.80', '0.85', '0.90', '0.95'))
    # # 显示colorbar
    # cbar = plt.colorbar(gci)
    # cbar.set_ticks(np.linspace(0, 10))
    # # cbar.set_ticklabels(('160', '180', '200', '220', '240', '260', '280', '300'))
    # plt.show()

    plt.imshow(data)
    plt.colorbar()
    plt.show()

def save_fig(data, filename):
    plt.imshow(data)
    # plt.colorbar()
    plt.savefig(filename, dpi=120)
    # plt.show()


def show_map(lon, lat):
    plt.figure()
    m = Basemap(projection='mill', lat_ts=10, llcrnrlon=lon.min(), urcrnrlon=lon.max(),
                llcrnrlat=lat.min(), urcrnrlat=lat.max(), resolution='c')
    m.drawcoastlines(linewidth=0.25)
    m.drawcountries(linewidth=0.25)
    m.fillcontinents(color='coral', lake_color='aqua')
    m.drawmapboundary(fill_color='aqua')
    m.drawmeridians(np.arange(0, 360, 30))
    m.drawparallels(np.arange(-90, 90, 30))
    x, y = m(lon, lat)
    # cs = m.contour(x, y, data)
    plt.title('Geopotential Height Contour from Grib')
    plt.show()


if __name__ == '__main__':
    file = config.SURF_CMPA_FRT_5KM_10MIN + '/Z_SURF_C_BABJ_20171113001802_P_CMPA_FAST_CHN_0P05_10MIN-PRE-201711130010.GRB2'
    grb = read_file(file)
    print(grb['year'], grb['month'], grb['day'], grb['hour'], grb['minute'], grb['second'])

    # for x in grb._all_keys:
    #     print(x)
